"""
Docker容器终端WebSocket Consumer
基于现有SSH终端实现，提供容器终端功能
"""
import json
import asyncio
import logging
import threading
import queue
from channels.generic.websocket import AsyncWebsocketConsumer
from channels.db import database_sync_to_async
from django.contrib.auth.models import AnonymousUser
from .models import DockerServer
from .services import DockerService

logger = logging.getLogger(__name__)


class DockerTerminalConsumer(AsyncWebsocketConsumer):
    """Docker容器终端WebSocket Consumer"""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.ssh_client = None
        self.ssh_channel = None
        self.server_id = None
        self.container_id = None
        self.docker_server = None
        self.user = None
        self.is_connected = False
        self.terminal_size = {'cols': 80, 'rows': 24}
        self.read_thread = None
        self.output_queue = queue.Queue()
        self.output_task = None

    async def connect(self):
        """WebSocket连接建立"""
        try:
            # 获取用户信息（从中间件设置）
            self.user = self.scope.get("user")

            if not self.user or not getattr(self.user, 'is_authenticated', False):
                logger.warning("Docker终端WebSocket连接被拒绝：用户未认证")
                await self.close(code=4001)
                return

            # 获取服务器ID和容器ID
            self.server_id = self.scope['url_route']['kwargs']['server_id']
            self.container_id = self.scope['url_route']['kwargs']['container_id']

            logger.info(f"Docker终端连接请求: 用户={self.user.username}, 服务器={self.server_id}, 容器={self.container_id}")

            # 接受WebSocket连接
            await self.accept()

            # 发送连接状态
            await self.send(text_data=json.dumps({
                'type': 'status',
                'status': 'websocket_connected',
                'message': 'WebSocket连接已建立'
            }))

        except Exception as e:
            logger.error(f"Docker终端WebSocket连接失败: {e}")
            await self.close(code=4000)

    async def disconnect(self, close_code):
        """WebSocket连接断开"""
        await self.cleanup_connection()

    async def receive(self, text_data):
        """接收WebSocket消息"""
        try:
            data = json.loads(text_data)
            message_type = data.get('type')

            if message_type == 'connect':
                # 建立容器终端连接
                terminal_size = data.get('terminal_size', {'cols': 80, 'rows': 24})
                self.terminal_size = terminal_size
                await self.establish_container_connection()

            elif message_type == 'input' and self.ssh_channel:
                # 发送用户输入
                input_data = data.get('data', '')
                await self.send_input_to_container(input_data)

            elif message_type == 'resize':
                # 调整终端大小
                cols = data.get('cols', 80)
                rows = data.get('rows', 24)
                await self.resize_terminal(cols, rows)

            elif message_type == 'disconnect':
                # 断开连接
                await self.cleanup_connection()

        except json.JSONDecodeError:
            logger.error("Docker终端接收到无效的JSON数据")
        except Exception as e:
            logger.error(f"Docker终端处理WebSocket消息失败: {e}")

    async def establish_container_connection(self):
        """建立容器终端连接"""
        try:
            # 获取Docker服务器信息
            self.docker_server = await database_sync_to_async(
                DockerServer.objects.get
            )(pk=self.server_id, is_active=True)

            # 建立SSH连接到Docker服务器
            await self.connect_to_docker_server()

            # 执行docker exec命令进入容器
            await self.exec_into_container()

        except Exception as e:
            logger.error(f"建立容器终端连接失败: {e}")
            await self.send(text_data=json.dumps({
                'type': 'error',
                'message': f'连接失败: {str(e)}'
            }))

    async def connect_to_docker_server(self):
        """连接到Docker服务器"""
        try:
            # 在线程中建立SSH连接
            def create_ssh_connection():
                try:
                    import paramiko
                    self.ssh_client = paramiko.SSHClient()
                    self.ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

                    self.ssh_client.connect(
                        hostname=self.docker_server.host,
                        port=self.docker_server.ssh_port,
                        username=self.docker_server.ssh_username,
                        password=self.docker_server.ssh_password,
                        timeout=10
                    )
                    return True
                except Exception as e:
                    logger.error(f"SSH连接失败: {e}")
                    return False

            # 在线程池中执行SSH连接
            loop = asyncio.get_event_loop()
            success = await loop.run_in_executor(None, create_ssh_connection)

            if not success:
                raise Exception("SSH连接失败")

        except Exception as e:
            logger.error(f"连接Docker服务器失败: {e}")
            raise

    async def exec_into_container(self):
        """执行docker exec进入容器"""
        try:
            def try_exec():
                try:
                    # 创建shell会话
                    channel = self.ssh_client.invoke_shell(
                        term='xterm-256color',
                        width=self.terminal_size['cols'],
                        height=self.terminal_size['rows']
                    )

                    # 通过SSH连接检测容器中的shell
                    available_shell = '/bin/sh'  # 默认使用sh

                    try:
                        # 使用SSH连接执行检测命令
                        stdin, stdout, stderr = self.ssh_client.exec_command(
                            f'docker exec {self.container_id} test -f /bin/bash',
                            timeout=5
                        )
                        exit_status = stdout.channel.recv_exit_status()

                        if exit_status == 0:
                            available_shell = '/bin/bash'
                            logger.info(f"容器 {self.container_id} 检测到bash，使用: {available_shell}")
                        else:
                            logger.info(f"容器 {self.container_id} 未检测到bash，使用默认: {available_shell}")

                    except Exception as e:
                        logger.warning(f"检测shell失败，使用默认shell /bin/sh: {e}")
                        available_shell = '/bin/sh'

                    logger.info(f"容器 {self.container_id} 最终使用shell: {available_shell}")

                    # 执行docker exec进入容器
                    docker_cmd = f'docker exec -it {self.container_id} {available_shell}'
                    logger.info(f"执行命令: {docker_cmd}")
                    channel.send(f'{docker_cmd}\n')

                    return channel

                except Exception as e:
                    logger.error(f"执行docker exec失败: {e}")
                    raise e

            # 在线程池中执行
            loop = asyncio.get_event_loop()
            self.ssh_channel = await loop.run_in_executor(None, try_exec)

            if not self.ssh_channel:
                raise Exception("无法创建容器终端会话")

            self.is_connected = True

            # 发送连接成功消息
            await self.send(text_data=json.dumps({
                'type': 'connected',
                'message': '容器终端连接成功'
            }))

            # 启动读取线程
            self.start_read_thread()

        except Exception as e:
            logger.error(f"进入容器失败: {e}")
            raise

    def start_read_thread(self):
        """启动读取线程"""
        def read_from_channel():
            try:
                while self.is_connected and self.ssh_channel:
                    if self.ssh_channel.recv_ready():
                        data = self.ssh_channel.recv(4096)
                        if data:
                            # 将数据放入队列
                            self.output_queue.put(data.decode('utf-8', errors='ignore'))
                    else:
                        import time
                        time.sleep(0.01)
            except Exception as e:
                logger.error(f"读取容器输出失败: {e}")

        self.read_thread = threading.Thread(target=read_from_channel, daemon=True)
        self.read_thread.start()

        # 启动输出处理任务
        self.output_task = asyncio.create_task(self.process_output_queue())

    async def process_output_queue(self):
        """处理输出队列"""
        try:
            while self.is_connected:
                try:
                    # 非阻塞地获取队列中的数据
                    data = self.output_queue.get_nowait()
                    await self.send_output(data)
                except queue.Empty:
                    # 队列为空，等待一下
                    await asyncio.sleep(0.01)
                except Exception as e:
                    logger.error(f"处理输出队列失败: {e}")
                    break
        except Exception as e:
            logger.error(f"输出队列处理任务失败: {e}")

    async def send_output(self, data):
        """发送输出到前端"""
        try:
            await self.send(text_data=json.dumps({
                'type': 'output',
                'data': data
            }))
        except Exception as e:
            logger.error(f"发送输出失败: {e}")

    async def send_input_to_container(self, data):
        """发送输入到容器"""
        try:
            if self.ssh_channel and self.is_connected:
                def send_data():
                    self.ssh_channel.send(data)

                loop = asyncio.get_event_loop()
                await loop.run_in_executor(None, send_data)
        except Exception as e:
            logger.error(f"发送输入到容器失败: {e}")

    async def resize_terminal(self, cols, rows):
        """调整终端大小"""
        try:
            self.terminal_size = {'cols': cols, 'rows': rows}

            if self.ssh_channel and self.is_connected:
                def resize():
                    self.ssh_channel.resize_pty(width=cols, height=rows)

                loop = asyncio.get_event_loop()
                await loop.run_in_executor(None, resize)
        except Exception as e:
            logger.error(f"调整终端大小失败: {e}")

    async def cleanup_connection(self):
        """清理连接"""
        try:
            self.is_connected = False

            # 取消输出处理任务
            if self.output_task and not self.output_task.done():
                self.output_task.cancel()
                try:
                    await self.output_task
                except asyncio.CancelledError:
                    pass

            if self.ssh_channel:
                self.ssh_channel.close()
                self.ssh_channel = None

            if self.ssh_client:
                self.ssh_client.close()
                self.ssh_client = None

            if self.read_thread and self.read_thread.is_alive():
                self.read_thread.join(timeout=1)

        except Exception as e:
            logger.error(f"清理连接失败: {e}")


